#!/usr/bin/python -u
# -*- coding: utf-8 -*-
# Author: Li Zhu < zhulipresent@gmail.com >

############ INPUT DATA ################
# You need POSCAR, POTCAR and INCAR.
# kgrids
kgrids = [0.045, 0.04, 0.035, 0.03, 0.025]

# M = 'M' for selects the original Monkhorst-Pack scheme,
# M = 'G' for generating meshes with their origin at the gamma  point.
M = 'M'

# vasp commond
cmd = 'mpiexec -n 12 /share/apps/vasp/vasp.4.6 > vasp.out'

# tolerence = 0.001 eV pre atom
tolerance = 0.001
############ END INPUT #################

import os
import math
def writekp(kgrid, M):
    def dot(x, y):
        return x[0]*y[0] + x[1]*y[1] + x[2]*y[2]

    def cross(x, y):
        z1 = x[1]*y[2] - x[2]*y[1]
        z2 = x[2]*y[0] - x[0]*y[2]
        z3 = x[0]*y[1] - x[1]*y[0]
        return [z1, z2, z3]
    def kmf(kgrid, gi):
        kd = int(gi/kgrid/2.0/math.pi)
        if kd == 0: kd = 1
        dd = gi/kd/2.0/math.pi
        if dd >= kgrid:
            for i in range(0, 10):
                kd += i
                dd = gi/kd/2.0/math.pi
                if dd <= kgrid: break
        return kd 
                    
    f = open('POSCAR')
    pp = []
    try:
        for line in f:
            pp.append(line.split())
    finally:
        f.close()
    l = []
    for item in pp[2:5]:
        l.append(map(float, item))
    # real Lattice Parameters
    ra = math.sqrt(l[0][0]**2 + l[0][1]**2 + l[0][2]**2)
    rb = math.sqrt(l[1][0]**2 + l[1][1]**2 + l[1][2]**2)
    rc = math.sqrt(l[2][0]**2 + l[2][1]**2 + l[2][2]**2)

    c = cross(l[1], l[2])
    volume = dot(l[0], c)
    g = []
    g1 = [ 2.0 * math.pi * item / volume for item in c]
    c = cross(l[2], l[0])
    g2 = [ 2.0 * math.pi * item / volume for item in c]
    c = cross(l[0], l[1])
    g3 = [ 2.0 * math.pi * item / volume for item in c]
    g = [g1, g2, g3]
    
    rl = []
    for i in range(0, 3):
        rl.append(math.sqrt(dot(g[i],g[i])))
    kmesh = []
    for i in range(0, 3):
        kmesh.append(kmf(kgrid, rl[i]))
    f = open('KPOINTS', 'w')
    f.write('A\n0\n')
    f.write(M + '\n')
    f.write('%2d %2d %2d\n' % tuple(kmesh))
    f.write('%2d %2d %2d\n' % (0,0,0))
    # debug
    #print 'volume: ', volume
    #print 'rl ', rl
    #print 'kp ', kmesh
    # end debug
    return kmesh

def getnatom():
    try:
        f = open('POSCAR')
    except:
        print 'ERROR: NO POSCAR'
        exit(0)
    fp = []
    try:
        for line in f:
            fp.append(line.split())
    finally:
        f.close()

    try:
        natom = sum(map(int, fp[5]))
    except:
        natom = sum(map(int, fp[6]))
    return natom

def run(kgrids, M, cmd, tolerance):
    natom = getnatom()
    energys = []
    kmeshs = []
    diffs = []
    last = 0.
    step = 0
    print " N       ENERGY       DIFF  KGRID"
    fe = open('energy.dat', 'w', 0)
    fd = open('diff.dat', 'w', 0)
    for kgrid in kgrids:
        kmesh = writekp(kgrid, M)
        kmeshs.append(kmesh)
        os.system(cmd)
        os.system('cp OUTCAR OUTCAR.' + str(kgrid))
        os.system('cp OSZICAR OSZICAR.' + str(kgrid))
        try:
            e = float(os.popen("grep free\ \  OUTCAR | awk '{print $5}' ").read())
        except:
            print "ERROR: can not get energy. (kgrid %5.3f)" % (kgrid)
            return 0
        energy = e / natom
        print >>fe, "%2d  %11.5f  %5.3f" % (step, energy, kgrid)
        diff = energy - last
        last = energy
        energys.append(energy)
        diffs.append(diff)
        # print log
        print "%2d  %11.5f  %9.5f  %5.3f" % (step, energy, diff, kgrid)
        if len(energys) > 1:
            print >>fd, "%2d  %9.5f" % (step, diff)
            if abs(diff) < tolerance:
                fe.close()
                fd.close()
                print "The difference is %7.5f ." % (diff)
                print "Kmesh: ( %3d %3d %3d ) is OK." % tuple(kmeshs[step - 1]) 
                print "Kgrid: %5.3f " % (kgrids[step - 1])
                return 0
        step += 1

    print 'NOT  CONVERGENCE'
    fe.close()
    fd.close()

if __name__ == '__main__':
    run(kgrids, M, cmd, tolerance)
